import os
import struct
import numpy as np
import matplotlib.pyplot as plt
 
def load_mnist(path, kind="train"):
	# label与image数据存储路径
    labels_path = os.path.join(path, '%s-labels.idx1-ubyte' % kind)
    images_path = os.path.join(path, '%s-images.idx3-ubyte' % kind)
 
    with open(labels_path, 'rb') as lbpath:
        # 由于label数据中前两个32位整型是数据描述而不是数据，所以要先读取出来，防止影响数据读取
        # '>'表示是big-endian大端模式存储，'I'表示一个无符号整数，所以'>II'就读取出数据描述
        magic, n = struct.unpack('>II', lbpath.read(8)) 
        # 读取数据部分
        labels = np.fromfile(lbpath, dtype=np.uint8)
 
    with open(images_path, 'rb') as imgpath:
        # 同上，读取数据描述
        magic, num, rows, cols =  struct.unpack('>IIII', imgpath.read(16)) 
        #每个图像的784个元素按照列向量排列，需要reshape才能使每一行为一个完整的图像
        images = np.fromfile(imgpath, dtype=np.uint8).reshape(len(labels), 784)
    return images, labels

#函数应用方式
path = "D:\\prj\\py\\ml\\data\\mnist"
X_train, y_train = load_mnist(path, kind="train")
X_test, y_test = load_mnist(path, kind="t10k")

# 如果想要获取label为i的第j张训练图像数据，可按照如下方法调用函数，注意此时获取的是大小为784的列向量
img = X_test[y_test == 9][1]
# 若想将以上图像显示出来，需要将该列向量reshape为28*28的二维图像数据
img = img.reshape(28, 28)

plt.imshow(img,cmap = plt.cm.gray)                                        
plt.show()  